import cv2
import os
import os.path as osp
import numpy as np
from PIL import Image


def joint_frame(lq_dir, hq_dir, save_dir):

    os.makedirs(save_dir, exist_ok=True)
    lq_file_list = sorted(os.listdir(lq_dir))
    hq_file_list = sorted(os.listdir(hq_dir))

    for lq_name, hq_name in zip(lq_file_list, hq_file_list):
        lq_path = osp.join(lq_dir, lq_name)
        hq_path = osp.join(hq_dir, hq_name)
        sv_path = osp.join(save_dir, hq_name)

        lq_img = Image.open(lq_path)
        hq_img = Image.open(hq_path)

        lq_ary = np.array(lq_img)
        hq_ary = np.array(hq_img)
        h, w, c = hq_ary.shape
        gap = np.ones((h, 30, c), dtype=hq_ary.dtype) * 255

        rst_ary = np.concatenate((lq_ary, gap, hq_ary), axis=1)

        rst_img = Image.fromarray(rst_ary)
        rst_img.save(sv_path)
        print(f'{sv_path} saved.')


def frame_to_video(frame_dir, save_path, fps=10):
    im_list = sorted(os.listdir(frame_dir))
    im_path_list = sorted([osp.join(frame_dir, im_name) for im_name in im_list])
    img = Image.open(im_path_list[0])
    img_size = img.size  #获得图片分辨率，im_dir文件夹下的图片分辨率需要一致

    fourcc = cv2.VideoWriter_fourcc(*'XVID')  #opencv版本是3
    videoWriter = cv2.VideoWriter(save_path, fourcc, fps, img_size)
    print('Start writing video.')
    for im_path in im_path_list:
        frame = cv2.imread(im_path)
        videoWriter.write(frame)
    videoWriter.release()
    print('Done.')


if __name__ == '__main__':
    lq_dir = '/home/mwh/frames/bicubic/5'
    hq_dir = '/home/mwh/frames/sfmn/5'
    save_dir = '/home/mwh/frames/results/5'
    video_path = '/home/mwh/frames/5.avi'

    joint_frame(lq_dir, hq_dir, save_dir)
    frame_to_video(save_dir, video_path)
